from torch.utils.data import Dataset
import os
from PIL import Image
from torchvision import transforms as T


class CrackDataSet(Dataset):
    def __init__(self, config, dataset_choose="train"):
        super(CrackDataSet, self).__init__()
        root_path = os.path.abspath(os.path.dirname(os.path.dirname(__file__)))
        if dataset_choose.upper() == "TRAIN":
            self.img_path = os.path.join(root_path, config["train_path"])
        elif dataset_choose.upper() == "TEST":
            self.img_path = os.path.join(root_path, config["test_path"])
        else:
            raise Exception("Illegal dataset name")
        self.img_names = os.listdir(self.img_path)
        self.dataset_choose = dataset_choose
        self.img_size = config["out_img_size"]

        normalize = T.Normalize(mean=[0.485], std=[0.229])
        if self.dataset_choose.upper() == "TEST":
            self.transfroms = T.Compose([
                T.Resize(self.img_size),
                T.CenterCrop(self.img_size),
                T.ToTensor(),
                normalize
            ])
        elif self.dataset_choose.upper() == "TRAIN":
            self.transfroms = T.Compose([
                T.Resize(self.img_size + self.img_size // 10),
                T.RandomResizedCrop(self.img_size),
                T.RandomHorizontalFlip(),  # 随机水平镜像
                T.ToTensor(),
                normalize
            ])

    def __len__(self):
        return len(self.img_names)

    def __getitem__(self, item):
        img_name = self.img_names[item]
        img_path = os.path.join(self.img_path, img_name)
        img = Image.open(img_path)
        img = self.transfroms(img)
        if img_name.split("_")[0] == "crack":
            lable = 0
        elif img_name.split("_")[0] == "normal":
            lable = 1
        return img, lable

if __name__ == '__main__':
    from config.Load_DefaultConfig import DefaultConfig
    config = DefaultConfig()
    temp = CrackDataSet(config.data_set)
    print(temp[5])